#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
下载管理对话框模块
提供模型下载管理界面
"""

import os
import sys
from pathlib import Path
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QListWidget, QProgressBar, QMessageBox
)
from PyQt5.QtCore import Qt, pyqtSignal, QThread

# 导入项目模块
sys.path.append(str(Path(__file__).resolve().parent.parent))
from core.model_manager import ModelManager
from config.paths import GDRIVE_FOLDER_ID, create_directories


class ModelDownloadThread(QThread):
    """模型下载线程"""
    
    progress = pyqtSignal(str)
    finished = pyqtSignal(bool)
    
    def __init__(self, model_manager):
        """
        初始化下载线程
        
        Args:
            model_manager: 模型管理器实例
        """
        super().__init__()
        self.model_manager = model_manager
    
    def run(self):
        """线程主函数"""
        try:
            self.progress.emit("开始下载模型文件...")
            success = self.model_manager.download_dependencies(force=True)
            
            if success:
                self.progress.emit("模型下载完成!")
                self.finished.emit(True)
            else:
                self.progress.emit("模型下载失败")
                self.finished.emit(False)
                
        except Exception as e:
            self.progress.emit(f"下载出错: {str(e)}")
            self.finished.emit(False)


class ModelDownloadDialog(QDialog):
    """模型下载管理对话框"""
    
    def __init__(self, parent=None):
        """初始化对话框"""
        super().__init__(parent)
        
        # 确保必要的目录结构存在
        create_directories()
        
        # 设置对话框属性
        self.setWindowTitle("模型下载管理")
        self.setMinimumSize(500, 350)
        
        # 创建模型管理器
        self.model_manager = ModelManager(ui_callback=self.log_message)
        
        # 初始化界面
        self.init_ui()
        
        # 检查模型状态
        self.check_model_status()
    
    def init_ui(self):
        """初始化用户界面"""
        # 主布局
        layout = QVBoxLayout(self)
        
        # 添加说明标签
        info_label = QLabel("下载Deep Ponies TTS所需的模型文件")
        layout.addWidget(info_label)
        
        # 模型状态
        self.status_label = QLabel("正在检查模型状态...")
        layout.addWidget(self.status_label)
        
        # 下载按钮
        self.download_button = QPushButton("下载模型文件")
        self.download_button.clicked.connect(self.download_models)
        layout.addWidget(self.download_button)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.progress_bar.setVisible(False)
        layout.addWidget(self.progress_bar)
        
        # 日志区域
        self.log_area = QListWidget()
        layout.addWidget(self.log_area)
        
        # 按钮区域
        button_layout = QHBoxLayout()
        
        self.refresh_button = QPushButton("刷新")
        self.refresh_button.clicked.connect(self.check_model_status)
        button_layout.addWidget(self.refresh_button)
        
        self.close_button = QPushButton("关闭")
        self.close_button.clicked.connect(self.accept)
        button_layout.addWidget(self.close_button)
        
        layout.addLayout(button_layout)
    
    def check_model_status(self):
        """检查模型状态"""
        self.log_message("正在检查模型文件...")
        
        if self.model_manager.check_model_files():
            self.status_label.setText("所有模型文件已下载")
            self.download_button.setEnabled(False)
            self.log_message("所有模型文件已存在，无需下载")
        else:
            self.status_label.setText("缺少必要的模型文件")
            self.download_button.setEnabled(True)
            self.log_message("检测到缺少模型文件，请点击下载按钮")
    
    def log_message(self, message, is_error=False):
        """
        记录消息到日志区域
        
        Args:
            message: 消息文本
            is_error: 是否为错误消息
        """
        prefix = "[错误] " if is_error else ""
        self.log_area.addItem(f"{prefix}{message}")
        self.log_area.scrollToBottom()
    
    def download_models(self):
        """下载模型文件"""
        self.log_message("准备下载模型文件...")
        self.download_button.setEnabled(False)
        self.refresh_button.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # 启动下载线程
        self.download_thread = ModelDownloadThread(self.model_manager)
        self.download_thread.progress.connect(self.log_message)
        self.download_thread.finished.connect(self.on_download_finished)
        self.download_thread.start()
        
        # 进度条动画
        self.progress_bar.setRange(0, 0)  # 设置为忙碌状态
    
    def on_download_finished(self, success):
        """
        下载完成的回调
        
        Args:
            success: 是否成功
        """
        self.progress_bar.setRange(0, 100)  # 恢复范围
        
        if success:
            self.progress_bar.setValue(100)
            self.status_label.setText("所有模型文件已下载")
            QMessageBox.information(self, "下载完成", "模型文件下载完成！")
        else:
            self.progress_bar.setValue(0)
            self.status_label.setText("模型下载失败")
            QMessageBox.warning(self, "下载失败", "模型文件下载失败，请检查网络连接或手动下载。")
        
        self.refresh_button.setEnabled(True)
        self.check_model_status() 